#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <netdb.h>
#include <errno.h>
#include <netinet/in.h>

#define BUFFER_SIZE 1024
#define PORT 1080
#define BACKLOG 10

// 函数声明
void *handle_client(void *arg);
int handle_socks5(int client_sock);
int resolve_hostname(const char *hostname, char *ip_address, size_t ip_len);
int handle_socks5_connect(int client_sock, unsigned char *buffer);

int main() {
    int server_sock, *client_sock;
    struct sockaddr_in server_addr, client_addr;
    socklen_t client_len = sizeof(client_addr);

    // 1. 创建套接字
    server_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (server_sock < 0) {
        perror("创建套接字失败");
        exit(EXIT_FAILURE);
    }

    // 2. 绑定套接字到本地地址和端口
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);
    server_addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(server_sock, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("绑定套接字失败");
        close(server_sock);
        exit(EXIT_FAILURE);
    }

    // 3. 启动监听
    if (listen(server_sock, BACKLOG) < 0) {
        perror("监听失败");
        close(server_sock);
        exit(EXIT_FAILURE);
    }

    printf("SOCKS代理服务器正在监听端口 %d...\n", PORT);

    // 接受连接
    while (1) {
        client_sock = malloc(sizeof(int));
        *client_sock = accept(server_sock, (struct sockaddr *)&client_addr, &client_len);
        if (*client_sock < 0) {
            perror("接受连接失败");
            free(client_sock);
            continue;
        }

        // 4. 创建线程处理每个客户端
        pthread_t thread;
        if (pthread_create(&thread, NULL, handle_client, (void *)client_sock) != 0) {
            perror("创建线程失败");
            free(client_sock);
            continue;
        }
        pthread_detach(thread); // 自动释放线程资源
    }

    close(server_sock);
    return 0;
}

// 线程函数，处理每个客户端
void *handle_client(void *arg) {
    int client_sock = *(int *)arg;
    free(arg);

    unsigned char buffer[BUFFER_SIZE];
    ssize_t received;

    // 5. 读取第一个字节，判断 SOCKS 协议版本
    received = recv(client_sock, buffer, sizeof(buffer), 0);
    if (received <= 0) {
        perror("从客户端读取数据失败");
        close(client_sock);
        return NULL;
    }

    if (buffer[0] == 0x05) { // SOCKS5
        printf("处理 SOCKS5 请求...\n");
        if (handle_socks5(client_sock) < 0) {
            printf("处理 SOCKS5 请求失败。\n");
        }
    } else if (buffer[0] == 0x04) { // SOCKS4（此实现未支持）
        printf("SOCKS4 尚未完全支持。\n");
    } else {
        printf("未知的 SOCKS 协议版本: %d\n", buffer[0]);
    }

    close(client_sock);
    return NULL;
}

// 处理 SOCKS5 请求
int handle_socks5(int client_sock) {
    unsigned char buffer[BUFFER_SIZE];
    ssize_t received;

    // 读取客户端支持的认证方法
    received = recv(client_sock, buffer, sizeof(buffer), 0);
    if (received <= 0) {
        perror("读取认证方法失败");
        return -1;
    }

    // 响应 "无需认证"
    unsigned char response[2] = {0x05, 0x00};
    if (send(client_sock, response, sizeof(response), 0) < 0) {
        perror("发送认证响应失败");
        return -1;
    }

    // 读取客户端请求
    received = recv(client_sock, buffer, sizeof(buffer), 0);
    if (received <= 0) {
        perror("读取客户端请求失败");
        return -1;
    }

    if (buffer[1] == 0x01) { // CONNECT 命令
        return handle_socks5_connect(client_sock, buffer);
    } else {
        printf("不支持的命令: %d\n", buffer[1]);
        return -1;
    }
}

// 处理 SOCKS5 CONNECT 命令
int handle_socks5_connect(int client_sock, unsigned char *buffer) {
    char target_address[256] = {0};
    uint16_t target_port;
    int remote_sock;
    struct sockaddr_in remote_addr;

    // 根据地址类型提取目标地址和端口
    if (buffer[3] == 0x01) { // IPv4
        inet_ntop(AF_INET, &buffer[4], target_address, sizeof(target_address));
        target_port = (buffer[8] << 8) | buffer[9];
    } else if (buffer[3] == 0x03) { // 域名
        unsigned char domain_length = buffer[4];
        strncpy(target_address, (char *)&buffer[5], domain_length);
        target_address[domain_length] = '\0';
        target_port = (buffer[5 + domain_length] << 8) | buffer[6 + domain_length];
    } else {
        printf("不支持的地址类型: %d\n", buffer[3]);
        return -1;
    }

    printf("正在连接到 %s:%d...\n", target_address, target_port);

    // 如果是域名，解析为 IP
    if (buffer[3] == 0x03) {
        char ip_address[INET_ADDRSTRLEN];
        if (resolve_hostname(target_address, ip_address, sizeof(ip_address)) < 0) {
            printf("解析域名失败: %s\n", target_address);
            return -1;
        }
        strncpy(target_address, ip_address, sizeof(target_address));
    }

    // 创建套接字连接目标服务器
    remote_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (remote_sock < 0) {
        perror("创建远程套接字失败");
        return -1;
    }

    remote_addr.sin_family = AF_INET;
    remote_addr.sin_port = htons(target_port);
    inet_pton(AF_INET, target_address, &remote_addr.sin_addr);

    if (connect(remote_sock, (struct sockaddr *)&remote_addr, sizeof(remote_addr)) < 0) {
        perror("连接目标服务器失败");
        close(remote_sock);
        return -1;
    }

    printf("成功连接到 %s:%d\n", target_address, target_port);

    // 向客户端发送连接成功的响应
    unsigned char response[10] = {0x05, 0x00, 0x00, 0x01};
    memset(&response[4], 0, 6); // BIND 地址和端口（此处未使用）
    send(client_sock, response, sizeof(response), 0);

    // 在客户端和目标服务器之间转发数据
    fd_set fds;
    int max_fd = (client_sock > remote_sock ? client_sock : remote_sock) + 1;

    while (1) {
        FD_ZERO(&fds);
        FD_SET(client_sock, &fds);
        FD_SET(remote_sock, &fds);

        if (select(max_fd, &fds, NULL, NULL, NULL) < 0) {
            perror("select 调用失败");
            break;
        }

        if (FD_ISSET(client_sock, &fds)) {
            ssize_t received = recv(client_sock, buffer, sizeof(buffer), 0);
            if (received <= 0) break;
            send(remote_sock, buffer, received, 0);
        }

        if (FD_ISSET(remote_sock, &fds)) {
            ssize_t received = recv(remote_sock, buffer, sizeof(buffer), 0);
            if (received <= 0) break;
            send(client_sock, buffer, received, 0);
        }
    }

    close(remote_sock);
    return 0;
}

// 解析主机名为 IP 地址
int resolve_hostname(const char *hostname, char *ip_address, size_t ip_len) {
    struct addrinfo hints, *res;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_INET; // IPv4

    if (getaddrinfo(hostname, NULL, &hints, &res) != 0) {
        perror("解析主机名失败");
        return -1;
    }

    struct sockaddr_in *addr = (struct sockaddr_in *)res->ai_addr;
    inet_ntop(AF_INET, &(addr->sin_addr), ip_address, ip_len);

    freeaddrinfo(res);
    return 0;
}
